import lcm
import numpy as np
from my_lcm import Request, Response
import torch

class LCM_Action_Publish:
    def __init__(self):
        self.lc = lcm.LCM()
        self.subscription = self.lc.subscribe("LCM_OBS", self.handle_Message)

        #观测序列
        self.obs = np.zeros(45, dtype=np.float32)
        
        #控制命令
        self.command = [1.0, 0.0, 0.0]

        #角速度缩放因子
        self.omega_scale = 0.25

        #控制命令缩放因子
        self.cmd_scale_x = 2.0
        self.cmd_scale_y = 2.0
        self.cmd_scale_z = 0.25

        #关节位置缩放因子
        self.dof_pos_scale = 1.0
        self.dof_vel_scale = 0.05

        #action初始化
        self.action = [0.0,0.0,0.0,  0.0,0.0,0.0,  0.0,0.0,0.0,  0.0,0.0,0.0]

        #关节角度初始化
        self.default_angles = [0.0,-1.0,2.0,0.0,-1.0,2.0,0.0,-1.0,2.0,0.0,-1.0,2.0]

        #加载策略
        self.policy_path = "/home/saw/C++/sim2sim/model/policy_5.pt"
        self.policy = torch.jit.load(self.policy_path)

    #发送action
    def send_action(self, data):
        policy_output = Response() #获取Response

        policy_output.q_exp = self.action

        self.lc.publish("LCM_ACTION", policy_output.encode())


    def handle_Message(self, channel, data):
        msg = Request.decode(data)

        #观测序列赋值
        self.obs[0] = msg.omega[0]*self.omega_scale
        self.obs[1] = msg.omega[1]*self.omega_scale
        self.obs[2] = msg.omega[2]*self.omega_scale
        self.obs[3:6] = msg.gravity_orientation

        self.obs[6] = self.command[0] * self.cmd_scale_x
        self.obs[7] = self.command[1] * self.cmd_scale_y
        self.obs[8] = self.command[2] * self.cmd_scale_z

        for i in range(12):
            self.obs[9 + i] = (msg.q[i] - self.default_angles[i]) * self.dof_pos_scale
            self.obs[9 + 12 + i] = msg.dq[i] * self.dof_vel_scale  

        self.obs[9 + 2 * 12 : 9 + 3 * 12] = self.action


        #模型推理
        obs_tensor = torch.from_numpy(self.obs).unsqueeze(0)
        self.action = self.policy(obs_tensor).detach().numpy().squeeze()


    def run(self):
        self.lc.handle()
 

if __name__ == "__main__":
    client = LCM_Action_Publish()
    while True:
        client.send_action("hello")
        client.run()